{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Surface restructuring analysis of LAMMPS NEB trajectory:\n",
    "## Preparation of LAMMPS/ASE optimization"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Header"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "import os\n",
    "import numpy as np\n",
    "import matplotlib.pyplot as plt\n",
    "\n",
    "print('Please cite DOI: 10.1021/acs.jpcc.9b04863.')\n",
    "print('This script analyzes LAMMPS NEB calculations of interlayer surface restructuring events (as prepared via MD2NEB.py) and prepares images for LAMMPS/ASE optimization of each event.')\n",
    "print('This script is interactive and requires user input. Please read carefully before proceeding:\\n')\n",
    "\n",
    "print('Note 1: Intended only for interlayer surface restructuring events in FCC metals.')\n",
    "print('Note 2: Intended only for periodic slab models with vacuum along the z-direction.')\n",
    "print('Note 3: Intended only for LAMMPS/ASE optimizations (ionic relaxation for IS/FS; dimer method for TS).')\n",
    "\n",
    "print('Note 4: Requires a LAMMPS DATA file containing the equilibrated box dimensions.')\n",
    "print('Note 5: Requires a LAMMPS DATA file containing the 0K box dimensions.')\n",
    "print('Note 6: Requires a file listing the event numbers chosen to be optimized, line-by-line.')\n",
    "print('Note 7: Requires NEB images in the form of LAMMPS DUMP file, named \"coords.final.$i\", 0 <= $i <= Nimg-1.')\n",
    "print('Note 8: Requires energies evaluated separately via Gaussian process (gp_neb.py), if using GP force field.\\n')"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Extract unit cell & system information"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "pwd = os.getcwd()\n",
    "\n",
    "# Load LAMMPS data file after NPT equilibration\n",
    "name = input('Enter the name of the LAMMPS DATA file containing the equilibrated box dimensions : ')\n",
    "file_eq = pwd + '/' + name\n",
    "log = open(file_eq,'r')\n",
    "log_lines = log.readlines()\n",
    "log_lines = [line.split() for line in log_lines]\n",
    "\n",
    "for line_index, line in enumerate(log_lines):\n",
    "    if line:\n",
    "        for l in line:\n",
    "            \n",
    "            if l == 'xlo':\n",
    "                xlo = float(line[0])\n",
    "                xhi = float(line[1])\n",
    "                \n",
    "            elif l == 'ylo':\n",
    "                ylo = float(line[0])\n",
    "                yhi = float(line[1])\n",
    "            \n",
    "            elif l == 'zlo':\n",
    "                zlo = float(line[0])\n",
    "                zhi = float(line[1])\n",
    "            \n",
    "            elif l == 'xy':\n",
    "                xy = float(line[0])\n",
    "                xz = float(line[1])\n",
    "                yz = float(line[2])\n",
    "\n",
    "lat_eq = np.array([[xhi-xlo, 0.0, 0.0], [xy, yhi-ylo, 0.0], [xz, yz, zhi-zlo]])    # LAMMPS NVT lattice matrix\n",
    "\n",
    "log.close()\n",
    "\n",
    "# Load LAMMPS data file after 0K equilibration\n",
    "name = input('Enter the name of the LAMMPS DATA file containing the 0K box dimensions : ')\n",
    "file_0K = pwd + '/' + name\n",
    "log = open(file_0K,'r')\n",
    "log_lines = log.readlines()\n",
    "log_lines = [line.split() for line in log_lines]\n",
    "\n",
    "for line_index, line in enumerate(log_lines):\n",
    "    if line:\n",
    "        for l in line:\n",
    "            \n",
    "            if l == 'xlo':\n",
    "                xlo = float(line[0])\n",
    "                xhi = float(line[1])\n",
    "                \n",
    "            elif l == 'ylo':\n",
    "                ylo = float(line[0])\n",
    "                yhi = float(line[1])\n",
    "            \n",
    "            elif l == 'zlo':\n",
    "                zlo = float(line[0])\n",
    "                zhi = float(line[1])\n",
    "            \n",
    "            elif l == 'xy':\n",
    "                xy = float(line[0])\n",
    "                xz = float(line[1])\n",
    "                yz = float(line[2])\n",
    "\n",
    "lat_0K = np.array([[xhi-xlo, 0.0, 0.0], [xy, yhi-ylo, 0.0], [xz, yz, zhi-zlo]])    # LAMMPS 0K lattice matrix\n",
    "\n",
    "xlo_bound = xlo + np.min(np.array([0.0, xy, xz, xy+xz]))    # Box bounds required for LAMMPS DUMP file\n",
    "xhi_bound = xhi + np.max(np.array([0.0, xy, xz, xy+xz]))\n",
    "ylo_bound = ylo + np.min(np.array([0.0, yz]))\n",
    "yhi_bound = yhi + np.max(np.array([0.0, yz]))\n",
    "zlo_bound = zlo\n",
    "zhi_bound = zhi\n",
    "\n",
    "bound_0K = np.array([[xlo_bound, xhi_bound, xy], [ylo_bound, yhi_bound, xz], [zlo_bound, zhi_bound, yz]])\n",
    "\n",
    "log.close()\n",
    "\n",
    "GP = input('Were the energies evaluated separately via Gaussian process (gp_neb.py)? Yes or No : ')\n",
    "\n",
    "print('Unit cell & system information extracted.\\n')"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## List of events to be optimized"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# List of events chosen to be optimized after visual inspection of the NEB results\n",
    "name = input('Enter the name of the file listing the event numbers chosen to be optimized, line-by-line : \\n Warning: All specified events must visually contain a single, well-defined event. Please quit now and perform visual inspection if not so.\\n')\n",
    "file_ev = pwd + '/' + name\n",
    "ev = open(file_ev,'r')\n",
    "ev_lines = ev.readlines()\n",
    "ev_lines = [line.split() for line in ev_lines]\n",
    "\n",
    "ls = []\n",
    "for line in ev_lines:\n",
    "    ls.append(line[0])\n",
    "    \n",
    "ev.close()\n",
    "\n",
    "print('\\n')"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Extract IS, TS, FS - LAMMPS DUMP file preparation"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# IS = initial state\n",
    "# TS = transition state\n",
    "# FS = final state\n",
    "\n",
    "print('Preparing LAMMPS DUMP files of IS, TS, FS in respective folders under ./ev_***/ for each event...')\n",
    "\n",
    "for l in ls:\n",
    "    \n",
    "    path = pwd + '/ev_' + str(l)\n",
    "    \n",
    "    if GP == 'Yes':\n",
    "        \n",
    "        file_mep = path + '/gp_neb.txt'    # List of image number & energy calculated via Gaussian process\n",
    "        mep = open(file_mep,'r')\n",
    "        mep_lines = mep.readlines()\n",
    "        mep_lines = [line.split() for line in mep_lines]\n",
    "        mep_lines[0] = np.array(mep_lines[0]).astype(int)\n",
    "        mep_lines[1] = np.array(mep_lines[1]).astype(float)\n",
    "        \n",
    "        f = plt.figure()    # Save MEP plot as a PDF file\n",
    "        plt.plot(mep_lines, 'k.', markersize=10)\n",
    "        plt.plot(mep_lines, 'k--', markersize=1)\n",
    "        plt.ylabel('Energy (eV)')\n",
    "        plt.xlabel('Image number')\n",
    "        file_pdf = path1 + '/mep.pdf'\n",
    "        f.savefig(file_pdf)\n",
    "\n",
    "    else:\n",
    "        \n",
    "        file_log = path + '/log.lammps'    # List of NEB distance & energy of the images\n",
    "        log = open(file_log,'r')\n",
    "        log_lines = log.readlines()\n",
    "        log_lines = np.array(log_lines)\n",
    "        N_line = log_lines.shape[0]\n",
    "        log_lines = [line.split() for line in log_lines]\n",
    "\n",
    "        file_mep = path + '/mep.txt'    # MEP = minimum energy pathway\n",
    "        mep = open(file_mep,'w')\n",
    "\n",
    "        os.chdir(path)\n",
    "        N_img = os.popen('ls -1 coords.final.* | wc -l').readlines()    # Extract number of images\n",
    "        N_img = [line.split() for line in N_img]\n",
    "        N_img = np.array(N_img).astype(int)\n",
    "        N_img = int(N_img)\n",
    "\n",
    "        for i in range(N_img,0,-1):    # Extract the converged NEB energy landscape (last entry)\n",
    "            out = log_lines[N_line-1][-2*i] + '\\t' + log_lines[N_line-1][-2*i+1] + '\\n'\n",
    "            mep.write(out)\n",
    "\n",
    "        mep.close()\n",
    "        log.close()\n",
    "    \n",
    "        mep = open(file_mep,'r')\n",
    "        mep_lines = mep.readlines()\n",
    "        mep_lines = [line.split() for line in mep_lines]\n",
    "        mep_lines = np.array(mep_lines).astype(float)        \n",
    "\n",
    "        f = plt.figure()    # Save MEP plot as a PDF file\n",
    "        plt.plot(mep_lines, 'k.', markersize=10)\n",
    "        plt.plot(mep_lines, 'k--', markersize=1)\n",
    "        plt.ylabel('Energy (eV)')\n",
    "        plt.xlabel('NEB distance')\n",
    "        file_pdf = path + '/mep.pdf'\n",
    "        f.savefig(file_pdf)\n",
    "    \n",
    "    # Scan MEP for IS, TS, FS\n",
    "    \n",
    "    for line_index, line in enumerate(mep_lines):    # TS = highest-energy image\n",
    "        if line[1] == np.nanmax(mep_lines[:,1]):\n",
    "            TS = line_index\n",
    "        \n",
    "    for i in range(TS,0,-1):    # Start at TS and climb down to the left to find IS\n",
    "        if mep_lines[i,1] - mep_lines[i-1,1] > 0:\n",
    "            if i > 1:\n",
    "                continue\n",
    "            else:\n",
    "                IS = i-1\n",
    "        else:\n",
    "            IS = i\n",
    "            break\n",
    "    \n",
    "    for i in range(TS,N_img-1):    # Start at TS and climb down to the right to find FS\n",
    "        if mep_lines[i,1] - mep_lines[i+1,1] > 0:\n",
    "            if i < N_img-2:\n",
    "                continue\n",
    "            else:\n",
    "                FS = i+1\n",
    "        else:\n",
    "            FS = i\n",
    "            break\n",
    "    \n",
    "    file_Nimg = path + '/Nimg.txt'    # List of image numbers corresponding to IS, TS, FS\n",
    "    Nimg = open(file_Nimg,'w')\n",
    "    \n",
    "    for i_index, i in enumerate([IS, TS, FS]):\n",
    "        for j_index, j in enumerate(['IS','TS','FS']):\n",
    "            if i_index == j_index:\n",
    "                os.mkdir(path + '/' + j)    # Create separate directories for IS, TS, FS\n",
    "                \n",
    "                out = j + '\\t' + str(i) + '\\n'\n",
    "                Nimg.write(out)\n",
    "                \n",
    "                file_dump = path + '/' + j + '/' + j + '.dmp'    # LAMMPS DUMP files for optimization\n",
    "                dump = open(file_dump,'w')\n",
    "        \n",
    "        file_img = path + '/coords.final.' + str(i)    # LAMMPS DUMP files of converged NEB images\n",
    "        img = open(file_img,'r')\n",
    "        img_lines = img.readlines()\n",
    "        \n",
    "        for line_index, line in enumerate(img_lines):\n",
    "            if (line_index < 5) or (line_index == 8):    # Write header\n",
    "                dump.write(line)\n",
    "            elif 5 <= line_index <= 7:    # Write 0K box bounds\n",
    "                out = '%f %f %f\\n'%(bound_0K[line_index-5][0], bound_0K[line_index-5][1], bound_0K[line_index-5][2])\n",
    "                dump.write(out)\n",
    "        \n",
    "        img_lines = [line.split() for line in img_lines]\n",
    "        \n",
    "        for line_index, line in enumerate(img_lines):\n",
    "            if line_index > 8:\n",
    "                \n",
    "                line[0:2] = np.array(line[0:2]).astype(int)    # ID, type\n",
    "                line[2:] = np.array(line[2:]).astype(float)    # x, y, z\n",
    "                \n",
    "                ID = line[0]\n",
    "                Type = line[1]\n",
    "                x = line[2]\n",
    "                y = line[3]\n",
    "                z = line[4]\n",
    "                r = np.array([x, y, z])    # Cartesian coordinates\n",
    "                Rep = np.dot(r, np.linalg.inv(lat_eq))    # Direct coordinates\n",
    "                r_0K = np.dot(Rep, lat_0K)    # Cartesian coordinates rescaled to 0K lattice\n",
    "                \n",
    "                out = '%i %i %f %f %f\\n'%(ID, Type, r_0K[0], r_0K[1], r_0K[2])    # Write rescaled entries\n",
    "                dump.write(out)\n",
    "        \n",
    "        img.close()\n",
    "        dump.close()\n",
    "    \n",
    "    Nimg.close()\n",
    "    mep.close()\n",
    "\n",
    "print('LAMMPS DUMP files generated > ./ev_***/{IS/,TS/,FS/}')"
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python 3",
   "language": "python",
   "name": "python3"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 3
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython3",
   "version": "3.7.3"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 2
}
